from tasks.handler import *
from django.conf import settings
from django.core.management import BaseCommand
from tasks.helper import get_loop_redis, friend_loads_json

import asyncio, logging

RunTimeLimit = 30

TaskRegister = {
    DemoTask.job_name: handle_demo
}


class Command(BaseCommand):
    """ 执行任务 """
    queue = None

    async def do_task(self, task: dict):
        try:
            await TaskRegister[task['job_name']](task)
        except Exception as e:
            logging.exception('do_task', exc_info=e)

        await self.queue.get()

    async def handle_tasks(self, redis):
        while True:
            await self.queue.put(0)
            key_task = await redis.blpop(settings.REDIS_TASK_TAG, timeout=10)
            if not key_task or key_task[0].decode() != settings.TASK_TAG: continue
            task = friend_loads_json(key_task[1])
            asyncio.ensure_future(self.do_task(task))

    def cancel_task(self):
        for task in asyncio.Task.all_tasks():
            task.cancel()

    def handle(self, *args, **options):
        loop, redis = get_loop_redis()
        self.queue = asyncio.Queue(RunTimeLimit)
        loop.run_until_complete(self.handle_tasks(redis))
        self.cancel_task()
